import tushare as ts
from datetime import time
from datetime import date, datetime, timedelta
from vnpy.trader.constant import Status, Interval, Direction
import pandas as pd

from vnpy_ctastrategy import (
    CtaTemplate,
    StopOrder,
    TickData,
    BarData,
    TradeData,
    OrderData,
    BarGenerator,
    ArrayManager,
)

cpi_path = 'E:\\GitCloneProject\\quant25\\01Src\\vnpy\\environment_evaluation\\data\\cpi_data.csv'

class MacroMarketStrategy(CtaTemplate):
    """宏观经济环境+股市环境三因子复合策略（双均线+布林带+ATR RSI）"""
    author = "Youpeng Hu"

    # === 参数设置 ===
    fast_window = 5        # 快速均线周期
    slow_window = 10       # 慢速均线周期
    boll_window = 20       # 布林带窗口
    boll_dev = 1.8         # 标准差倍数
    rsi_window = 14        # RSI周期
    rsi_entry = 30         # RSI入场阈值
    atr_window = 14        # ATR周期
    atr_ma_window = 10     # ATR均线窗口
    stop_multiplier = 3.0  # 止损系数
    max_pos = 1000         # 最大持仓
    risk_ratio = 0.02      # 风险比例
    dividend_days_before = 2  # 分红前N天清仓
    cpi_threshold = 0.5  # CPI阈值

    # 数据存储
    cpi_data = {}
    last_cpi_month = None
    current_month = None

    parameters = [
        "fast_window", "slow_window",
        "boll_window", "boll_dev",
        "rsi_window", "rsi_entry",
        "atr_window", "atr_ma_window",
        "stop_multiplier", "max_pos", "risk_ratio",
        "dividend_days_before", "cpi_threshold"
    ]

    variables = [
        "fast_ma", "slow_ma",
        "boll_up", "boll_mid", "boll_down",
        "rsi_value", "atr_value", "atr_ma",
        "intra_trade_low", "cpi_condition_met"
    ]

    def __init__(self, cta_engine, strategy_name, vt_symbol, setting):
        super().__init__(cta_engine, strategy_name, vt_symbol, setting)
        
        # K线生成器
        self.bg = BarGenerator(
            self.on_bar,
            interval=Interval.DAILY,
            daily_end=time(15, 0)
        )
        self.am = ArrayManager(size=100)
        self.last_trade_date = None
        self.pro = ts.pro_api(token='600713cc777ee2505142620a527eebe2c27a973b5fe89fce9fe8f07f')
        # 存储分红日期数据
        self.dividend_dates = []
        self.load_dividend_data()  # 初始化时加载数据
        self.load_cpi_data()  # 加载CPI数据

    def load_dividend_data(self):
        """根据动态vt_symbol加载分红数据"""
        # 1. 解析vt_symbol为ts_code
        try:
            symbol_part, exchange_suffix = self.vt_symbol.split(".")
        except ValueError:
            self.write_log(f"无效vt_symbol格式：{self.vt_symbol}，应为'代码.交易所'格式（如600848.SH）")
            return
        
        # 2. 转换为Tushare需要的ts_code格式
        ts_code = f"{symbol_part}.{exchange_suffix}"
        
        # 3. 调用Tushare接口
        try:
            raw_df = self.pro.dividend(
                ts_code=ts_code, 
                fields="ex_date,div_proc,stk_div"
            )
        except Exception as e:
            self.write_log(f"股票{ts_code}分红数据获取失败：{str(e)}")
            return

        # 4. 处理返回的DataFrame
        self.dividend_dates = []
        if not raw_df.empty:
            for _, row in raw_df.iterrows():
                item = row.to_dict()

    def is_near_dividend(self, bar_date: date) -> bool:
        """判断是否在分红敏感期"""
        for ex_date in self.dividend_dates:
            # 计算关键日期
            record_date = ex_date - timedelta(days=1)   
            danger_start = record_date - timedelta(days=self.dividend_days_before)
            
            if danger_start <= bar_date <= ex_date:
                return True
        return False

    def load_cpi_data(self):
        """加载本地CPI数据"""
        df = pd.read_csv(cpi_path, dtype={"month": str})
        self.cpi_data = df.set_index("month").to_dict()["nt_yoy"]
        print(df.head())

    def on_init(self):
        """策略初始化"""
        self.load_bar(30)  # 加载30根K线
        self.write_log("策略初始化完成")

    def on_bar(self, bar: BarData):
        """核心交易逻辑"""
        # 数据有效性校验
        if bar.close_price <= 0 or bar.datetime.date() == self.last_trade_date:
            return
            
        self.am.update_bar(bar)
        if not self.am.inited:
            return

        # 转换K线时间为date对象
        bar_date = bar.datetime.date()
        
        # 检查是否在分红敏感期
        if self.is_near_dividend(bar_date):
            self.write_log(f"进入分红敏感期，当前日期{bar_date}")
            
            # 强制清仓
            if self.pos != 0:
                self.sell(bar.close_price * 0.995, abs(self.pos))
                self.cancel_all()
                
            # 禁止新开仓
            return
        # === 指标计算 ===
        self._calculate_indicators(bar)
        
        # === 交易信号 ===
        # volume_condition是当前成交量大于过去5日均量的1.2倍。
		# l1_signal是快速均线大于慢速，且收盘价高于布林中轨；
		# l2_signal是ATR值大于其均线，且收盘价高于布林上轨；
		# l3_signal是RSI小于入场阈值，但当前值大于过去5个的最小值。
        
        # 检查CPI条件
        current_date = bar.datetime.strftime("%Y%m")
        self.cpi_condition_met = False

        if current_date in self.cpi_data:
            last_month = (datetime.strptime(current_date, "%Y%m") - pd.DateOffset(months=1)).strftime("%Y%m")
            print(last_month)
            if last_month in self.cpi_data:
                self.cpi_condition_met = self.cpi_data[last_month] < self.cpi_threshold

        volume_condition = bar.volume > self.am.volume[-5:].mean() * 1.2
        l1_signal = (self.fast_ma > self.slow_ma) and (bar.close_price > self.boll_mid)
        l2_signal = (self.atr_value > self.atr_ma) and (bar.close_price > self.boll_up)
        l3_signal = (self.rsi_value < self.rsi_entry) and (self.rsi_value > self.rsi_array[-5:].min())
        
        # print(l1_signal, l2_signal, l3_signal, self.cpi_condition_met)
        
        # === 交易执行 ===
        self.cancel_all()
        self.add_risk_control()  # 激活风控模块
        
        # 修正后的开仓条件
        if ((l1_signal and l2_signal) or (l2_signal and l3_signal)) \
           and self.pos == 0 \
           and volume_condition \
           and self.cpi_condition_met:
            
            size = self._calculate_position_size(bar)
            price = min(bar.close_price * 1.005, bar.high_price * 0.995)
            self.buy(price, size)
            self.intra_trade_low = bar.low_price
            self.last_trade_date = bar.datetime.date()

        # 持仓管理
        elif self.pos > 0:
            self._manage_position(bar)

    def _calculate_indicators(self, bar):
        """集中计算技术指标"""
        self.fast_ma = self.am.sma(self.fast_window)
        self.slow_ma = self.am.sma(self.slow_window)
        self.boll_up, self.boll_down = self.am.boll(self.boll_window, self.boll_dev)
        self.boll_mid = self.am.sma(self.boll_window)
        self.rsi_value = self.am.rsi(self.rsi_window)
        atr_array = self.am.atr(self.atr_window, array=True)
        self.atr_value = atr_array[-1]
        self.atr_ma = atr_array[-self.atr_ma_window:].mean()
        self.rsi_array = self.am.rsi(self.rsi_window, array=True)
        self.rsi_value = self.rsi_array[-1]  # 当前RSI值

    def _manage_position(self, bar: BarData):
        """持仓管理"""
        # 更新intra_trade_low为当前最低价和之前的intra_trade_low中的较小值，
        # 计算止损价和止盈价（布林中轨的1.05倍）。如果当前价格跌破止损或超过止盈，就平仓。
        self.intra_trade_low = min(self.intra_trade_low, bar.low_price)
        stop_price = self._calculate_stop_price(bar)
        take_profit = self.boll_mid * 1.05

        if bar.close_price < stop_price or bar.close_price > take_profit or not self.cpi_condition_met:
            self.sell(bar.close_price, abs(self.pos))
            self.last_trade_date = bar.datetime.date()  # 平仓日期记录

    def _calculate_stop_price(self, bar) -> float:
        """复合止损策略"""
        # 结合ATR止损（最近低点减去ATR乘以系数）和趋势止损（慢速均线的97%），取两者中的较大值作为止损价
        atr_stop = self.intra_trade_low - self.atr_value * self.stop_multiplier
        trend_stop = self.slow_ma * 0.97
        return max(atr_stop, trend_stop)

    def _calculate_position_size(self, bar) -> int:
        """基于波动率的仓位管理"""
        # 仓位计算_calculate_position_size是根据风险比例计算风险金额，然后除以ATR的两倍，再与最大持仓比较取小值
        risk_amount = self.cta_engine.capital * self.risk_ratio
        return min(int(risk_amount / (self.atr_value * 2)), self.max_pos)

    def add_risk_control(self):
        """增强风控模块"""
        # 持仓量风控
        # 检查持仓是否超过最大持仓的80%，如果是就停止策略
        if self.pos >= self.max_pos * 0.8:
            self.write_log("触发单品种持仓限制")
            self.on_stop()
        
        # 周回撤控制
        """
        if hasattr(self.cta_engine, 'get_weekly_drawdown'):
            if self.cta_engine.get_weekly_drawdown() > 0.10:
                self.sell(bar.close_price, abs(self.pos))
                self.write_log("触发周最大回撤限制")
        """
				
